import os
import django

os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'nexora.settings')
django.setup(set_prefix=False)

import asyncio
import re
import zoneinfo
from websockets.exceptions import ConnectionClosed
from websockets.frames import CloseCode
from websockets.asyncio.server import serve, ServerConnection
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework import HTTP_HEADER_ENCODING
from rest_framework.exceptions import APIException
from datetime import datetime
from services.sshclient import SSHClient
from tree.models import Host, Track
from django.conf import settings

# 或下行获取Model类Host
# from django.apps import apps as django_apps
# Host = django_apps.get_model('jumpserver', 'Host', require_ready=False)

pattern = re.compile("^/tree/webshell/(?P<host_id>[^/]+)")
logq = asyncio.Queue()


def get_host(host_id) -> dict:  # 同步函数，handler中使用时改为异步否则报错
    host = Host.objects.filter(is_deleted=False).get(pk=host_id)  # 查不到就抛异常
    conf = {
        "hostname": host.ip,
        "username": host.username,
        "port": 22,
    }
    if host.password:
        conf["password"] = host.password
    if host.ssh_pkey_path:
        full_key_path = settings.JUMP_UPLOADS_DIR / host.ssh_pkey_path
        conf["ssh_pkey_path"] = str(full_key_path)
    return conf


# 特别注意，这个函数执行应该是client.output有返回就send到浏览器端
async def get_remote_output(ws: ServerConnection, client: SSHClient):
    while not client.closed:  # 如果当前client关闭，就没必要执行这个task了，循环结束
        data = await asyncio.to_thread(client.output)  # 把client.output执行变异步
        if data is not None:
            await ws.send(data)
    await ws.close()  # 既然client已经不能从远程主机获取信息了，和浏览器的ws也可以断开了
    print("ws.close()@@@")


# websocket 是和前端浏览器的websocket通信
# client    是和远程主机的SSH通信

# 获取了消息，组织好内容，打入到队列中，启动一个Log协程单独处理
async def get_remote_message(client: SSHClient, user_id, host_id, source_ip: str):
    while True:
        data = await asyncio.to_thread(client.message)
        if data is not None:
            await logq.put(dict(
                user_id=user_id, host_id=host_id,
                source_ip=source_ip, **data))
        else:
            break



async def log():
    while True:
        if logq.empty():
            await asyncio.sleep(0.1)
            continue
        data: dict = await logq.get()
        print('|' * 30)
        # {'user_id': 1, 'host_id': '1', 'source_ip': '::1', 'type': 1, 'command': '', 'output': '', 'created': 1752161978.018835}
        print(data)
        try:
            t = Track()
            t.user_id = data['user_id']
            t.host_id = data['host_id']
            t.source_ip = data['source_ip']
            t.op_type = data['type']
            # 不带时区，会有naive警告
            d = datetime.fromtimestamp(data['created'], zoneinfo.ZoneInfo(settings.TIME_ZONE))
            t.op_date = d
            if d := data.get('command', ''):
                t.command = d
            if d := data.get('output', ''):
                t.op_result = d
            # t.save() # 入数据库
            await asyncio.to_thread(t.save)  # 转异步
        except Exception as e:
            print(e, "$$$")
        print('|' * 30)


async def handler(websocket: ServerConnection):
    print(type(websocket), websocket)
    # 认证，第一条从客户端发来的是token，去jwt认证
    raw_token = await websocket.recv()
    print(type(raw_token), raw_token)
    jwtauth = JWTAuthentication()
    code = CloseCode.NORMAL_CLOSURE
    reason = ""
    try:
        # 参考 JWTAuthentication authenticate方法
        validated_token = jwtauth.get_validated_token(raw_token.encode(HTTP_HEADER_ENCODING))
        user = await asyncio.to_thread(jwtauth.get_user, validated_token)
        print(type(user), user)
        if not user:
            code = CloseCode.INTERNAL_ERROR
            reason = "认证失败"

        # check path
        path = websocket.request.path  # websockets.http11.Request.path str
        print(f"path={path}")
        if m := pattern.match(path):
            pass  # 匹配成功，继续
        else:
            code = CloseCode.INTERNAL_ERROR
            reason = f"请求路径{path}不合法"
            return

        # 查询主机信息
        host_id = m.group("host_id")
        # config = get_host(host_id) # 同步操作，报错
        config = await asyncio.to_thread(get_host, host_id)
        print(config)
        print('@' * 30)

        # Echo Server
        hello = "{:%Y/%m/%d %H:%M:%S} 欢迎{}, 请开始输入命令\r\n".format(datetime.now(), user)
        await websocket.send(hello)
        with SSHClient(**config) as client:
            # 为当前SSH客户端获取消息协程
            asyncio.run_coroutine_threadsafe(
                get_remote_message(client, user.id, host_id, websocket.remote_address[0]),
                asyncio.get_running_loop()
            )
            # 要把get_remote_output异步函数塞到当前事件大循环中执行
            # run_coroutine_threadsafe提交协程到指定的事件循环中反复执行
            asyncio.run_coroutine_threadsafe(
                get_remote_output(websocket, client),
                asyncio.get_running_loop()
            )
            # await get_remote_output(websocket, client) # 不行，这是固定位置执行一次
            await asyncio.sleep(1)  # 停一下，让上面从服务器端拿回数据去显示
            while not client.closed:
                data = None
                # xterm attach后，网页中每输入一个字符就会发过来
                try:
                    data = await asyncio.wait_for(websocket.recv(), timeout=0.1)
                except asyncio.TimeoutError as e:
                    continue
                except ConnectionClosed:
                    # 客户端主动关闭连接（比如用户关了浏览器标签）
                    print("Client disconnected gracefully.")
                    break  # 退出循环，不再读取
                except Exception as ee:
                    # 上面recv时，如果websocket close了，就抛异常了
                    print(f"[{type(ee)},,{ee},,{data!r}!@#]")
                    raise ee  # 连接已经断了，往外抛异常
                print('&' * 30)
                print(type(data), data.encode())
                # await websocket.send(data) # 回显不要了
                # 把客户端发来的按键直接转发
                client.send(data)

    except APIException as e:
        print(e)
        code = CloseCode.GOING_AWAY
        reason = e.default_detail
    except Exception as e:
        print(type(e), e)
        code = CloseCode.GOING_AWAY
        reason = str(e)  # 关闭原因
    finally:
        print(code, reason, "fin")
        await websocket.close(code, reason)  # 关闭
        print('+' * 30)


async def main():
    asyncio.run_coroutine_threadsafe(
        log(),  # 审计信息入库
        asyncio.get_running_loop()
    )
    async with serve(handler, "0.0.0.0", 10008) as server:
        await server.serve_forever()


if __name__ == "__main__":
    asyncio.run(main())
